import time
from collections import defaultdict
from datetime import datetime
from itertools import islice

import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
from torch_geometric.data import Data

from GATModel import GATModel
from genBOM import genData
from genGraph import saveGraph


class PEVO:
    def __init__(self, bom_structure, bom_attr):
        self.bom_structure = bom_structure
        self.bom_attr = bom_attr
        self.nodes = set()
        self.edges = defaultdict(lambda: defaultdict(int))
        self.node_to_idx = {}
        self.edge_index = None
        self.edge_weight = None
        self.node_features = None

    # 遍历BOM结构，提取节点和边
    def extract_nodes_edges(self, bom_structure):
        for parent, children in bom_structure.items():
            if isinstance(parent, tuple):
                parent_name, _ = parent
            else:
                parent_name = parent
            self.nodes.add(parent_name)

            for child in children:
                if isinstance(child, tuple):
                    child_name, child_qty = child
                elif isinstance(child, str):
                    child_name = child
                    # 非元组，默认数量为1
                    child_qty = 1
                else:
                    continue
                self.nodes.add(child_name)
                self.edges[parent_name][child_name] += child_qty

                if isinstance(children, dict) and child in children:
                    # 嵌套情况下，递归获取
                    self.extract_nodes_edges({child: children[child]})

    def build_graph(self):

        # 开始构建有向图，记录开始时间
        build_graph_start_time = time.time()

        # 遍历BOM结构，提取节点和边
        self.extract_nodes_edges(self.bom_structure)

        # 节点映射到索引
        self.node_to_idx = {node: i for i, node in enumerate(self.nodes)}

        # 创建边索引和边权重
        edge_index_list = []
        edge_weight_list = []

        for src, dests in self.edges.items():
            for dest, weight in dests.items():
                edge_index_list.append([self.node_to_idx[src], self.node_to_idx[dest]])
                edge_weight_list.append(weight)

        self.edge_index = torch.tensor(edge_index_list, dtype=torch.long).t().contiguous()
        self.edge_weight = torch.tensor(edge_weight_list, dtype=torch.float)

        # 构建节点的特征向量
        # 字符串不能直接转换为tensor向量，这里通过One-hot编码转为0/1编码
        ######
        # 采用独热编码可能会导致维度比较高，考虑稀疏矩阵或者降维
        ######

        # 1. 先遍历，提取唯一属性
        unique_properties = set()
        for properties in self.bom_attr.values():
            unique_properties.update(properties.keys())
        # 2. 构建属性值到索引的映射字典
        property_to_idx = {prop: i for i, prop in enumerate(unique_properties)}
        # 3. 对属性信息进行独热编码
        node_features_list = []
        for node in self.nodes:
            features = self.bom_attr.get(node, {})
            one_hot_encoding = [0] * len(unique_properties)
            for prop, value in features.items():
                prop_idx = property_to_idx.get(prop)
                if prop_idx is not None:
                    one_hot_encoding[prop_idx] = 1
            feature_vector = torch.tensor(one_hot_encoding, dtype=torch.float)
            node_features_list.append(feature_vector)
        # 4. 属性集合表达为节点特征
        self.node_features = torch.stack(node_features_list)

        # 创建Data对象
        self.data = Data(x=self.node_features, edge_index=self.edge_index, edge_attr=self.edge_weight)

        print("pyG对象：", self.data)

        # 结束构建有向图，记录结束时间
        build_graph_end_time = time.time()
        print(f"构建有向图花费时间: {build_graph_end_time - build_graph_start_time} 秒")

    def train(self, epochs=1000):
        # 开始训练模型，记录开始时间
        train_start_time = time.time()
        # 实例化模型
        hidden_size = 8
        out_size = len(self.nodes)
        model = GATModel(self.node_features.size(1), hidden_size, out_size)

        # Train model
        # 添加weight_decay进行L2正则化
        optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=1e-5)
        model.train()
        losses = []
        for epoch in range(epochs):
            optimizer.zero_grad()
            out = model(self.data.x, self.data.edge_index)
            # 均方差通常用于回归任务
            # torch.rand随机数不适合评估模型性能，需要提供GT
            loss = F.mse_loss(out, torch.rand((len(self.nodes), out_size)))
            loss.backward()
            optimizer.step()
            losses.append(loss.item())
            if epoch % 10 == 0:
                print(f"Epoch {epoch}: Loss {loss.item()}")

        # 生成损失函数的变化图
        plt.figure()
        plt.plot(range(epochs), losses)
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.title('Loss over epochs')
        plt.savefig(f'loss_over_epochs_{datetime.now().strftime("%Y%m%d_%H%M%S")}.png', dpi=300, bbox_inches='tight')

        # Save model
        torch.save(model.state_dict(), "model.pth")

        # 结束训练模型，记录结束时间
        train_end_time = time.time()
        print(f"模型训练花费时间: {train_end_time - train_start_time} 秒")
        print("GNN Dimension : ", model)

    def test(self, componentName, model):

        model.load_state_dict(torch.load("model.pth"))
        model.eval()

        if componentName not in self.node_to_idx:
            return "Component not found."

        # 获取组件的索引
        component_idx = self.node_to_idx[componentName]

        # 使用图神经网络进行预测，需要将该节点及其相关的图结构信息构建成一个 Data 对象
        # 构建图数据
        data = Data(x=self.node_features, edge_index=self.edge_index)

        # 进行预测
        with torch.no_grad():
            predictions = model(data.x, data.edge_index)

        # 获取所有可能的子件及其预测分数
        scores = {}
        for idx, node in enumerate(self.nodes):
            if idx != component_idx:  # 排除自身
                scores[node] = predictions[component_idx, idx].item()

        # 排序并选择得分最高的子件
        recommended_components = sorted(scores.items(), key=lambda x: x[1], reverse=True)

        return recommended_components


# 使用BOMGraphBuilder类
if __name__ == '__main__':
    """
    bom_structure = {
        'WGPart1': {
            'SBB3': {('SBB8', 2), 'SBB4', 'SBB5'},
            'SBB4': {'SBB3': {('SBB8', 2)}, ('SBB9', 2): {}},
            'SBB5': {'SBB1', 'SBB6'},
            'SBB6': {'SBB5', ('SBB3', 10)},
            ('SBB8', 3): {('SBB1', 10)}
        },
        ('WGPart2', 2): {},
        'WGPart3': {'SBB1': {'SBB2'}},
        'WGPart4': {'SBB7': {'SBB8', ('SBB9', 2)}}
    }

    bom_attr = {
        "SBB1": {"大小": "10寸", "材质": "玻璃", "重量": "200g"},
        "SBB2": {"型号": "ABC123", "功耗": "50W"},
        "SBB3": {"品牌": "XYZ", "尺寸": "5mm"},
        "SBB4": {"类型": "机械键盘", "颜色": "黑色", "重量": "300g"},
        "SBB5": {"型号": "Intel Core i7", "功耗": "65W"}
    }
    """

    bom_structure, bom_attr = genData(part_nums=20, sbb_nums=100)
    saveGraph(bom_structure)

    pevo = PEVO(bom_structure, bom_attr)
    pevo.build_graph()
    pevo.train(epochs=1000)

    component_name = "SBB3"
    recommendations = pevo.test(component_name, GATModel(pevo.node_features.size(1), 8, len(pevo.nodes)))
    print(f"Recommendations for {component_name}:")
    # 获取前三个得分高的
    for component, score in islice(recommendations, 3):
        print(f"    {component}: {score:.4f}")
